#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Copyright    :   Jihua Lab 2021.
@File    :   robot_control_node.py
@Author  :   WANG Hao 
@Contact :   wanghao@jihualab.com、haowanghk@gmail.com(Github)
@License :   (C)Copyright 2020-2021, JHL
@Version :   2.4.1
@Create Time    :  2021-12-20
@Modify Time      @Author    @Version    @Desciption
-------------------        --- -------    ---  --------   --- -----------
@File Description :   DUAL_ARM ros robot control metapackage  
@Update: 2021-12-23 add dual_arm_control threadding and nodes, pre for testing
         2022-2-28 modified the  dual_arm_control threadding and nodes paramters
         2022-3-02 modified the ping host ip
         2022-03-08 update control system and structure
         2022-03-09 update servo info pub and state sub 

'''

# the import lib
import logging
import threading
import subprocess
import re

import uservo
import serial
import time

# import robot motion control pkg
from MotionGroup import MotionGroup
from ServoPose import Pose
from wrist_gripper import WristGripperControl
from robot_arm import ArmMotionControl

# inport ros py
import rospy
from sensor_msgs.msg import JointState
from joint_state_exchange.msg import RobotState
from dual_arm_base_control.msg import TelemetryInfo
from dual_arm_base_control.srv import CtrlModeSelect
from dual_arm_base_control.msg import NetworkState
from std_msgs.msg import String

# flag
exitFlag = 0  # threading control flag

# Global const para
HOST_STATE = True  # the homer host ip port state , default is true
SHORT_TIME_DELAY = 0.015  # time delay para
MEDIA_TIME_DELAY = 0.05
LONG_TIME_DELAY = 1.0

controlPeriod = 1000
leftWristAndGripperAngle = [0.0, 0.0]  # 存储手腕手爪开合临时变量
rightWristAndGripperAngle = [0.0, 0.0]

ServoModeHasDoneFlag = False

# publisher
SAMPLING_FREQUENCY = 5  # units: Hz
fileName = "robot_arm_records.txt"
POSE_NUM = rospy.get_param('/teach_mode/way_point_num')  # default pose number is 300

# PARAMS
g_ending_servo_idx = [0, 1, 2, 3]
Pi = 3.1415926535
Pi_Para = 57.29577951471995  # Pi_Para = 180.0/Pi

VEL = {'Very_low_spped': 10, 'LOW_SPEED': 15, 'MID_SPEED': 45, 'HIGH_SPEED': 60}

# @ brief:抓取动作初始位置，相对零位姿
#         手臂竖直状态,绝对角度控制模式
#: 双臂机器人起飞及降落状态时手臂必须处于预备蜷缩位置姿态 DUAL_ARM_INIT_POSE
DUAL_ARM_INIT_POSE = [140.0, 50.0, 72.0, 0,
                      -2.0, -80, 60.0, 0,
                      40, 40]
# 伺服模式下的零点位置，即手臂竖直向下（小臂关节角度为0）状态为0位，手腕90°、手爪10°
DUAL_ARM_ZERO_ANGLE_POSE = [140.0, 50.0, -30.0, 90,
                            -2.0, -80, -35, 90,
                            10, 10]
# 抓取目标位置
DUAL_ARM_TARGET_POSE = [140.0, 50.0, -30.0, 0,
                        2.0, -82, -35, 0,
                        80, 90]

# 计算关节角度绝对角度数值与关节角度相对于零点位置的相对转动角度
dual_arm_joint_angle_controller = []
dual_arm_joint_angle_relatively = []


# dual_arm_joint_angle_relatively = dual_arm_joint_angle_controller - DUAL_ARM_START_POSE


class RosRobotControl(object):
    def __init__(self):
        # global dual arm servo and wrist/gripper motor object

        # TODO: debug 2022-03-08
        self.dual_arm_node_handler = None
        self.wrist_gripper_node_handler = None
        self.servo_mode_has_done_flag = ServoModeHasDoneFlag
        self.leftWristAndGripperAngle = leftWristAndGripperAngle
        self.rightWristAndGripperAngle = rightWristAndGripperAngle
        # set the init paramters
        self._init_arm_servo()
        self.state_pub_freq = self.get_pub_freq()
        self.servo_list = self.get_servo_list()

        # set the rosnode and sub \ pub node handler
        self.robot_state_pub = None  # pub the robot arm state infomation
        self.network_state_sub = None  # sub the newwork state
        self.servo_sub = None  # sub the servo mode control cmd
        self.ros_node_init()

    def ros_run(self):
        """Run ros wrapper of ``rospy.spin()``."""
        self.spin()
        # rospy.spin()

    def ros_shutdown(self):
        """Shudown hook to close the node."""
        robot_control_obj.pose_init()        
        rospy.on_shutdown()
        
        rospy.loginfo('Shutting down ' + rospy.get_name() + ' node.')

    def ros_node_init(self):
        # def the ros node and the publihser and subscribers
        rospy.loginfo("-----ROS NODE TESTING ---- ")
        try:
            rospy.init_node(rospy.get_param('/node_info/robot_control_node_name'), anonymous=True)
            rospy.loginfo("-----Dual Arm Control and State Publsiher Node init success!-----")
        except rospy.ROSInitException as rosInitEx:
            rospy.logerr("-----Dual Arm Control and State Publsiher Node init ERROR:{}".format(str(rosInitEx)))

        # define the publisher
        try:
            # ros pub node init 
            # node name:   dual_arm_state_pub_node
            # node topic:  /servo_info
            self.robot_state_pub = rospy.Publisher(rospy.get_param('/node_info/state_pub_topic_name'),
                                                   TelemetryInfo,
                                                   queue_size=rospy.get_param('/node_info/state_pub_queue_size'))
            # pub_rate = rospy.Rate(self.state_pub_freq)         # default pub frequence is 5 
            rospy.loginfo("-----Robot State Publsiher Topic init success!-----")
        except rospy.ROSException:
            rospy.logerr("-----Robot State Publsiher Topic init ERROR !-----")

        # define subscriber
        self.network_state_sub = rospy.Subscriber(rospy.get_param('/node_info/network_pub_topic_name'),
                                                  NetworkState,
                                                  self.check_network_state,
                                                  queue_size=rospy.get_param('/node_info/network_pub_queue_size'))

        self.servo_sub = rospy.Subscriber('/robot_arm_state',
                                          RobotState,
                                          self.robot_control_callback,
                                          queue_size=5)

    def _init_arm_servo(self):
        # dual_arm class obj init
        dual_arm_ = ArmMotionControl(arm_name="dual_arm",
                                     dof=rospy.get_param('/arm_info/dual_arm_dof'),
                                     CONTROL_MODE="DAMPING",
                                     with_wrist=False,
                                     is_debug=True,
                                     damping_power=1000)
        wrist_gripper_ = WristGripperControl(servo_channel=g_ending_servo_idx, is_debug=True)

        # 初始化舵机管理器类对象、串口端口地址
        rospy.loginfo("----DUAL ARM STATE INIT !-----")
        dual_arm_.creat_uservo_obj(SERVO_BAUDRATE=115200)
        if dual_arm_.servo_buffer(is_scan_servo=True):
            # dual arm and gripper init
            dual_arm_servo = dual_arm_.get_arm_servo_list()
            dual_arm_.arm_state_init(force_angle_init=False)
            dual_arm_.set_all_servo_mode(CONTROL_MODE="CIRCLE")
            wrist_gripper_.all_ending_init(motion_group=4)
            dual_arm_.logging_arm_info()  # logging servo info
            dual_arm_.get_arm_current_angle()  # logging servo angle info
            rospy.loginfo("----DUAL ARM STATE INIT: control obj init success!-----")
            rospy.loginfo("----DUAL ARM STATE INIT : Servo Manager Started with servo {}".format(dual_arm_servo))
 
            time.sleep(1)
        else:
            rospy.logerr("ERROR: dual arm and gripper init failed!!!")

        # global control obj from dual_control_manager
        self.dual_arm_node_handler = dual_arm_
        self.wrist_gripper_node_handler = wrist_gripper_

        time.sleep(1)
        rospy.loginfo("-----Zero POSE TESTING END---- ")

    def get_servo_list(self):
        # get the servo list of robot arm
        output_servo_list = self.dual_arm_node_handler.get_arm_servo_list()
        return output_servo_list

    # get the robot state publisher frequence
    def get_pub_freq(self):
        freq = rospy.get_param('/node_info/state_pub_freq')
        return freq

    ##### take off and landing dual_arm pose init ####
    def pose_init(self):
        # servo index input check
        if (self.servo_list == [] or self.dual_arm_node_handler == []):
            rospy.logerr("ERROR INPUT PARAMTERS: servo_list {} ".format(self.servo_list))
            return False
        # set robot arms、wrist and grippers into landing pose
        for SERVO_ID in self.servo_list:
            self.dual_arm_node_handler.set_single_servo_pose(
                Pose=generate_pose_vtc(angle=DUAL_ARM_INIT_POSE[SERVO_ID],
                                       vel=VEL.get('LOW_SPEED'),
                                       iterval=2 * controlPeriod,
                                       t_acc=40,
                                       t_dec=40,
                                       power=9000,
                                       mean_dps=40,
                                       servo_id=SERVO_ID))
            # 左臂手腕手爪位置伺服，0（IO5）-》it-3，1(IO6)->it-[8]
            self.wrist_gripper_node_handler.set_servo_angle(servo_idx=0, angle=DUAL_ARM_INIT_POSE[3])
            self.wrist_gripper_node_handler.set_servo_angle(servo_idx=1, angle=DUAL_ARM_INIT_POSE[8])

            # 右臂手腕手爪位置伺服，2（IO12）-》it-7，3(IO13)->it-[9]
            self.wrist_gripper_node_handler.set_servo_angle(servo_idx=2, angle=DUAL_ARM_INIT_POSE[7])
            self.wrist_gripper_node_handler.set_servo_angle(servo_idx=3, angle=DUAL_ARM_INIT_POSE[9])
            rospy.loginfo("----Dual Arm and gripper Pose init success!! --- ")
            time.sleep(SHORT_TIME_DELAY)
        return True

    ####### dual arm robot auto mode ##########
    def start_servo_mode(self):
        # sub JointState and RobotState to control the servo motors 
        rospy.loginfo("---- THREADING  start servo mode !----")
        # control robot arm move to the zero position
        for SERVO_ID in self.servo_list:
            self.dual_arm_node_handler.set_single_servo_pose(
                Pose=generate_pose_vtc(angle=DUAL_ARM_ZERO_ANGLE_POSE[SERVO_ID],
                                       vel=VEL.get('VERY_LOW_SPEED'),
                                       iterval=2 * controlPeriod,
                                       t_acc=80,
                                       t_dec=60,
                                       power=9000,
                                       mean_dps=30,
                                       servo_id=SERVO_ID))
            time.sleep(SHORT_TIME_DELAY)

        # set the robot arm wrist joint and gripper angle
        # 左臂手腕手爪位置伺服，0（IO5）-》it-3，1(IO6)->it-[8]
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=0, angle=DUAL_ARM_ZERO_ANGLE_POSE[3])
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=1, angle=DUAL_ARM_ZERO_ANGLE_POSE[8])

        # 右臂手腕手爪位置伺服，2（IO12）-》it-7，3(IO13)->it-[9]
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=2, angle=DUAL_ARM_ZERO_ANGLE_POSE[7])
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=3, angle=DUAL_ARM_ZERO_ANGLE_POSE[9])
        time.sleep(1)

    def start_landing_mode(self):
        # force landing pose for dual_arm robot
        rospy.loginfo("----Start landing mode !----")
        servo_list = self.dual_arm_node_handler.get_arm_servo_list()
        for SERVO_ID in self.servo_list:
            self.dual_arm_node_handler.set_single_servo_pose(
                Pose=generate_pose_vtc(angle=DUAL_ARM_INIT_POSE[SERVO_ID],
                                       vel=VEL.get('VERY_LOW_SPEED'),
                                       iterval=2 * controlPeriod,
                                       t_acc=80,
                                       t_dec=60,
                                       power=9000,
                                       mean_dps=20,
                                       servo_id=SERVO_ID))
            time.sleep(SHORT_TIME_DELAY)

        # set the robot arm wrist joint and gripper angle
        # 左臂手腕手爪位置伺服，0（IO5）-》it-3，1(IO6)->it-[8]
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=0, angle=DUAL_ARM_INIT_POSE[3])
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=1, angle=DUAL_ARM_INIT_POSE[8])

        # 右臂手腕手爪位置伺服，2（IO12）-》it-7，3(IO13)->it-[9]
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=2, angle=DUAL_ARM_INIT_POSE[7])
        self.wrist_gripper_node_handler.set_servo_angle(servo_idx=3, angle=DUAL_ARM_INIT_POSE[9])

    ######### State PUB ##############
    def state_msg(self):
        # 创建舵机关节信息发布者话题
        # 默认状态发布频率5hz
        # main function to construct new control threading
        # 临时变量
        angle_tmp = []
        power_tmp = []
        tempture_tmp = []

        #### Create TelemetryInfo ####
        msg = TelemetryInfo()
        rospy.loginfo('-----Robot Telemetry Info Publishing-----')

        # publish the servo angle and other info
        for SERVO in self.servo_list:
            angle_tmp.append(self.dual_arm_node_handler.read_single_servo_angle(SERVO_ID=SERVO))
            power_tmp.append(self.dual_arm_node_handler.get_servo_status(SERVO_ID=SERVO).get('power'))
            tempture_tmp.append(self.dual_arm_node_handler.get_servo_status(SERVO_ID=SERVO).get('temp'))
            # rospy.loginfo(angle_tmp)

        #### Create TelemetryInfo msg ####
        msg.header.stamp = rospy.Time.now()
        msg.leftArmJointAngle.append(angle_tmp[0])
        msg.leftArmJointAngle.append(angle_tmp[1])
        msg.leftArmJointAngle.append(angle_tmp[2])

        msg.rightArmJointAngle.append(angle_tmp[3])
        msg.rightArmJointAngle.append(angle_tmp[4])
        msg.rightArmJointAngle.append(angle_tmp[5])

        msg.leftArmJointPower.append(power_tmp[0])
        msg.leftArmJointPower.append(power_tmp[1])
        msg.leftArmJointPower.append(power_tmp[2])
        msg.leftArmJointPower.append(0)
        msg.leftArmJointPower.append(0)

        msg.rightArmJointPower.append(power_tmp[3])
        msg.rightArmJointPower.append(power_tmp[4])
        msg.rightArmJointPower.append(power_tmp[5])
        msg.rightArmJointPower.append(0)
        msg.rightArmJointPower.append(0)

        msg.leftArmJointTemp.append(tempture_tmp[0])
        msg.leftArmJointTemp.append(tempture_tmp[1])
        msg.leftArmJointTemp.append(tempture_tmp[2])
        msg.leftArmJointTemp.append(0)
        msg.leftArmJointTemp.append(0)

        msg.rightArmJointTemp.append(tempture_tmp[3])
        msg.rightArmJointTemp.append(tempture_tmp[4])
        msg.rightArmJointTemp.append(tempture_tmp[5])
        msg.rightArmJointTemp.append(0)
        msg.rightArmJointTemp.append(0)

        return msg

    def rad2deg(self, rad_val):
        # 弧度制转角度值
        deg = (rad_val * Pi_Para)
        return deg

        ###########

    # check host state call back
    # get the HOST_STATE = False or True
    ############    
    def check_network_state(self, NetworkState):
        # check the target host ip state 
        # return Boolean type True or False
        if NetworkState.state:
            HOST_STATE = True
        else:
            HOST_STATE = False

    ##### function callback #####
    def robot_control_callback(self, RobotState):
        # check the network state
        if not HOST_STATE:
            rospy.loginfo(
                "----Robot Ros Controller Callback: HOST_STATE false- > robot arm move to landing pose !! --- ")
            self.pose_init()  # HOST_STATE false- > robot arm move to landing pose
            time.sleep(2)
        elif RobotState.init_flag:
            if not self.servo_mode_has_done_flag:
                # HOST_STATE is true and then init_flag is true for robot arm ready to servo ZERO ANGLE POSE
                # flag is true for robot servo mode is ready
                rospy.loginfo(
                    "----Robot Ros Controller Callback: CMD flag is true for robot servo mode is ready !! --- ")
                self.start_servo_mode()
                time.sleep(2)
                self.servo_mode_has_done_flag = True  # success switch to servo mode

            elif RobotState.cmd_flag:
                # cmd_flag : sub the control cmd and control robot arm and gripper
                if self.servo_list is None:
                    rospy.logerr("ERROR from (dual_arm_nh or wrist_gripper or servo_list) is None!")
                for SERVO_ID in self.servo_list[:3]:
                    # 0,1,2 left
                    if SERVO_ID == 1:
                        # 1号舵机转动方向相反
                        self.dual_arm_node_handler.set_single_servo_pose(
                            Pose=generate_pose_vtc(
                                angle=(DUAL_ARM_ZERO_ANGLE_POSE[SERVO_ID] - RobotState.leftArmPosition[
                                    SERVO_ID] * Pi_Para),
                                vel=VEL.get('LOW_SPEED'),
                                iterval=controlPeriod,
                                t_acc=80,
                                t_dec=60,
                                power=10000,
                                mean_dps=30,
                                servo_id=SERVO_ID))
                    else:
                        # servo 0,2 转动方向相同
                        self.dual_arm_node_handler.set_single_servo_pose(
                            Pose=generate_pose_vtc(
                                angle=(DUAL_ARM_ZERO_ANGLE_POSE[SERVO_ID] + RobotState.leftArmPosition[
                                    SERVO_ID] * Pi_Para),
                                vel=VEL.get('LOW_SPEED'),
                                iterval=1 * controlPeriod,
                                t_acc=80,
                                t_dec=60,
                                power=10000,
                                mean_dps=30,
                                servo_id=SERVO_ID))
                    time.sleep(SHORT_TIME_DELAY)

                for SERVO_ID in self.servo_list[3:7]:
                    # 4,5,6 right arm
                    self.dual_arm_node_handler.set_single_servo_pose(
                        Pose=generate_pose_vtc(
                            angle=(DUAL_ARM_ZERO_ANGLE_POSE[SERVO_ID] + RobotState.rightArmPosition[
                                SERVO_ID - 4] * Pi_Para),
                            vel=VEL.get('LOW_SPEED'),
                            iterval=1 * controlPeriod,
                            t_acc=80,
                            t_dec=60,
                            power=10000,
                            mean_dps=30,
                            servo_id=SERVO_ID))
                    time.sleep(SHORT_TIME_DELAY)
                # get the msg data, and exchange the unit from rad to deg using fast para Pi_Para
                # set the robot arm wrist joint and gripper angle
                # 左臂手腕手爪位置伺服，0（IO5）-》it-3，1(IO6)->it-[8]
                self.wrist_gripper_node_handler.set_servo_angle(servo_idx=0,
                                                                angle=RobotState.leftArmPosition[3] * Pi_Para)
                self.wrist_gripper_node_handler.set_servo_angle(servo_idx=1,
                                                                angle=RobotState.leftArmPosition[4] * Pi_Para)

                # 右臂手腕手爪位置伺服，2（IO12）-》it-7，3(IO13)->it-[9]
                self.wrist_gripper_node_handler.set_servo_angle(servo_idx=2,
                                                                angle=RobotState.rightArmPosition[3] * Pi_Para)
                self.wrist_gripper_node_handler.set_servo_angle(servo_idx=3,
                                                                angle=RobotState.rightArmPosition[4] * Pi_Para)
                # get the temp wrist and gripper angle value in degree
                self.leftWristAndGripperAngle[0] = RobotState.leftArmPosition[3] * Pi_Para
                self.leftWristAndGripperAngle[1] = RobotState.leftArmPosition[4] * Pi_Para

                self.rightWristAndGripperAngle[0] = RobotState.rightArmPosition[3] * Pi_Para
                self.rightWristAndGripperAngle[1] = RobotState.rightArmPosition[4] * Pi_Para

                time.sleep(SHORT_TIME_DELAY)
            else:
                rospy.loginfo(
                    "----Robot Ros Controller Callback: RobotState.cmd_flag is {}".format(RobotState.cmd_flag))
        else:
            rospy.loginfo(
                "----Robot Ros Controller Callback: Warning !!! Robot NOT ready !! RobotState.init_flag {}. --- ".format(
                    RobotState.init_flag))

    #############################################################
    def spin(self):
        rate = rospy.Rate(self.state_pub_freq)
        time_now = rospy.Time.now()
        while not rospy.is_shutdown():
            self.spinOnce()
            rate.sleep()

    #############################################################
    def spinOnce(self):
        # read robot arm state and return the generated msg
        output_state_msg = self.state_msg()
        output_state_msg.leftArmJointAngle.append(
            self.leftWristAndGripperAngle[0])  # append the left arm wrist angle
        output_state_msg.leftArmJointAngle.append(
            self.leftWristAndGripperAngle[1])  # append the left arm gripper angle
        output_state_msg.rightArmJointAngle.append(
            self.rightWristAndGripperAngle[0])  # append the right arm wrist angle
        output_state_msg.rightArmJointAngle.append(
            self.rightWristAndGripperAngle[1])  # append the right arm gripper angle
        try:
            self.robot_state_pub.publish(output_state_msg)
        except ValueError as ve:
            rospy.logwarn(str(ve))

def generate_pose_vtc(angle, servo_id, vel=100, iterval=5000, t_acc=35, t_dec=35, power=10000, mean_dps=70.0):
    # 增加关节舵机角度的vel和iterval参数动态调整接口API
    pose_tmp = {'angle': angle, 'is_mturn': False, 'interval': iterval,
                'vel': vel, 't_acc': t_acc,
                't_dec': t_dec, 'power': power,
                'mean_dps': mean_dps}
    gen_pose = Pose(POSE=pose_tmp, SERVO_ID=servo_id)
    return gen_pose


def main():
    # init the ropy nodelet and paramters defianation
    is_debug = rospy.get_param('/general/is_debug')
    ######### MAIN ##############
    rospy.loginfo("######### MAIN ##############!")
    robot_control_obj = RosRobotControl()
    try:
        robot_control_obj.ros_run()
    except KeyboardInterrupt:
        robot_control_obj.ros_shutdown()
        
        rospy.loginfo("-----退出主线程，主控节点结束!-----")


# main function
if __name__ == '__main__':
    # 创建一个logger_debug
    import sys

    # logger = logging.getLogger()
    # logger.setLevel(logging.DEBUG)  # Log等级总开关-debug
    # print("Logging Set Info: logger.setLevel(logging.DEBUG")
    try:
        main()
    except KeyboardInterrupt:
        rospy.loginfo("-----KeyboardInterrupt: 主程序正在退出!-----")
